import torch

# 初始化模型和优化器
model = torch.nn.TransformerEncoderLayer(1024, 16, 4096)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)

# 指定批次大小
batch_size = 2 # 你需要根据你的GPU容量调整这个值

# 确保CUDA可用
if torch.cuda.is_available():
    model = model.cuda()

# 初始化记录最大显存使用量
max_memory_used = 0

for i in range(3):
    # 创建随机输入数据并将其转移到GPU
    inputs = torch.rand([batch_size, 512, 1024], requires_grad=True)
    if torch.cuda.is_available():
        inputs = inputs.cuda()

    # 前向传播
    output = model(inputs)
    
    # 反向传播
    output.sum().backward()
    
    # 更新优化器
    optimizer.step()
    
    # 清除旧的梯度
    optimizer.zero_grad()
    
    # 检查和记录当前步骤后的最大显存使用
    current_memory = torch.cuda.max_memory_allocated()
    max_memory_used = max(max_memory_used, current_memory)
    
    # 可选：在每次迭代后重置最大内存跟踪器
    torch.cuda.reset_peak_memory_stats()

# 打印记录的最大显存使用量
print(max_memory_used)
print(f"Maximum GPU memory used: {max_memory_used / (1024**3):.15f} GB")

